# @Time   : 2020/10/6
# @Author : Shanlei Mu
# @Email  : slmu@ruc.edu.cn

"""
recbole.quick_start
########################
"""
import logging
from logging import getLogger

from recbole.config import Config
from recbole.data import create_dataset, data_preparation, data_preparation_head_tail
from recbole.utils import init_logger, get_model, get_trainer, init_seed
from recbole.utils.utils import set_color
import sys

def run_recbole(model=None, dataset=None, config_file_list=None, config_dict=None, saved=True):
    r""" A fast running api, which includes the complete process of
    training and testing a model on a specified dataset

    Args:
        model (str): model name
        dataset (str): dataset name
        config_file_list (list): config files used to modify experiment parameters
        config_dict (dict): parameters dictionary used to modify experiment parameters
        saved (bool): whether to save the model
    """
    # configurations initialization
    config = Config(model=model, dataset=dataset, config_file_list=config_file_list, config_dict=config_dict)
    init_seed(config['seed'], config['reproducibility'])

    # logger initialization
    init_logger(config)
    logger = getLogger()
    logger.info(sys.argv)
    logger.info(config)

    # dataset filtering
    dataset = create_dataset(config)
    logger.info(dataset)

    # dataset splitting
    # train_data, valid_data, test_data, head_data, tail_data = data_preparation_head_tail(config, dataset)
    train_data, valid_data, test_data= data_preparation(config, dataset)

    # model loading and initialization
    model = get_model(config['model'])(config, train_data).to(config['device'])
    logger.info(model)

    # trainer loading and initialization
    trainer = get_trainer(config['MODEL_TYPE'], config['model'])(config, model)

    # model training
    best_valid_score, best_valid_result = trainer.fit(
        train_data, valid_data, saved=saved, show_progress=config['show_progress']
    )

    # import numpy as np
    # import seaborn as sns
    # import matplotlib.pyplot as plt
    # from sklearn.decomposition import TruncatedSVD

    # embedding_matrix = model.item_embedding.weight[1:].cpu().detach().numpy()
    # svd = TruncatedSVD(n_components=2)
    # svd.fit(embedding_matrix)
    # comp_tr = np.transpose(svd.components_)
    # proj = np.dot(embedding_matrix, comp_tr)
    
    # cnt = {}
    # for i in dataset['item_id']:
    #     if i.item() in cnt:
    #         cnt[i.item()] += 1
    #     else:
    #         cnt[i.item()] = 1
    
    # freq = np.zeros(embedding_matrix.shape[0])
    # for i in cnt:
    #     freq[i-1] = cnt[i]
    
    # # freq /= freq.max()

    # sns.set(style='darkgrid')
    # sns.set_context("notebook", font_scale=1.8, rc={"lines.linewidth": 3, 'lines.markersize': 20})
    # plt.figure(figsize=(6, 4.5))
    # plt.scatter(proj[:, 0], proj[:, 1], s=1, c=freq, cmap='viridis_r')
    # plt.colorbar()
    # plt.xlim(-2, 2)
    # plt.ylim(-2, 2)
    # # plt.axis('square')
    # # plt.show()
    # plt.savefig(log_dir + '/' + config['model'] + '-' + config['dataset'] + '.pdf', format='pdf', transparent=False, bbox_inches='tight')
    
    # from scipy.linalg import svdvals
    # svs = svdvals(embedding_matrix)
    # svs /= svs.max()
    # np.save(log_dir + '/sv.npy', svs)

    # sns.set(style='darkgrid')
    # sns.set_context("notebook", font_scale=1.8, rc={"lines.linewidth": 3, 'lines.markersize': 20})
    # plt.figure(figsize=(6, 4.5))
    # plt.plot(svs)
    # # plt.show()
    # plt.savefig(log_dir + '/svs.pdf', format='pdf', transparent=False, bbox_inches='tight')

    # model evaluation
    test_result = trainer.evaluate(test_data, load_best_model=saved, show_progress=config['show_progress'])
    # head_test_result = trainer.evaluate(
    #     head_data, load_best_model=False, show_progress=True
    # )
    # tail_test_result = trainer.evaluate(
    #     tail_data, load_best_model=False, show_progress=True
    # )
    logger.info(set_color('best valid ', 'yellow') + f': {best_valid_result}')
    logger.info(set_color('test result', 'yellow') + f': {test_result}')
    # logger.info(set_color("head test", "yellow") + f": {head_test_result}")
    # logger.info(set_color("tail test", "yellow") + f": {tail_test_result}")

    return {
        'best_valid_score': best_valid_score,
        'valid_score_bigger': config['valid_metric_bigger'],
        'best_valid_result': best_valid_result,
        'test_result': test_result
    }


def objective_function(config_dict=None, config_file_list=None, saved=True):
    r""" The default objective_function used in HyperTuning

    Args:
        config_dict (dict): parameters dictionary used to modify experiment parameters
        config_file_list (list): config files used to modify experiment parameters
        saved (bool): whether to save the model
    """

    config = Config(config_dict=config_dict, config_file_list=config_file_list)
    init_seed(config['seed'], config['reproducibility'])
    logging.basicConfig(level=logging.ERROR)
    dataset = create_dataset(config)
    train_data, valid_data, test_data = data_preparation(config, dataset)
    model = get_model(config['model'])(config, train_data).to(config['device'])
    trainer = get_trainer(config['MODEL_TYPE'], config['model'])(config, model)
    best_valid_score, best_valid_result = trainer.fit(train_data, valid_data, verbose=False, saved=saved)
    test_result = trainer.evaluate(test_data, load_best_model=saved)

    return {
        'best_valid_score': best_valid_score,
        'valid_score_bigger': config['valid_metric_bigger'],
        'best_valid_result': best_valid_result,
        'test_result': test_result
    }

def load_data_and_model(model_file):
    r"""Load filtered dataset, split dataloaders and saved model.

    Args:
        model_file (str): The path of saved model file.

    Returns:
        tuple:
            - config (Config): An instance object of Config, which record parameter information in :attr:`model_file`.
            - model (AbstractRecommender): The model load from :attr:`model_file`.
            - dataset (Dataset): The filtered dataset.
            - train_data (AbstractDataLoader): The dataloader for training.
            - valid_data (AbstractDataLoader): The dataloader for validation.
            - test_data (AbstractDataLoader): The dataloader for testing.
    """
    import torch

    checkpoint = torch.load(model_file)
    config = checkpoint["config"]
    init_seed(config["seed"], config["reproducibility"])
    init_logger(config)
    logger = getLogger()
    logger.info(config)

    dataset = create_dataset(config)
    logger.info(dataset)
    train_data, valid_data, test_data = data_preparation(config, dataset)

    init_seed(config["seed"], config["reproducibility"])
    model = get_model(config["model"])(config, train_data.dataset).to(config["device"])
    model.load_state_dict(checkpoint["state_dict"])
    # model.load_other_parameter(checkpoint.get("other_parameter"))

    return config, model, dataset, train_data, valid_data, test_data
